Skip to content

refactor(tensor-pool): move behind a generic extension seam, opt-in and outside the 1.0 guarantees - #3152

Merged
trunk-io[bot] merged 3 commits into
mainfrom
pool-extract
Aug 13, 2026
Merged

refactor(tensor-pool): move behind a generic extension seam, opt-in and outside the 1.0 guarantees#3152
trunk-io[bot] merged 3 commits into
mainfrom
pool-extract

Conversation

@phil-opp

@phil-opp phil-opp commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Moves the pinned/CUDA tensor-pool transport (#2168, #2386, #2619) out of dora's core and behind a generic extension seam. It stays in the tree and stays usable — as an opt-in feature, explicitly outside the 1.0 compatibility guarantees.

Three commits, each independently reviewable:

  1. Remove the transport from dora's core — 3,048 lines from the Python binding, 973 from the daemon, five wire-protocol variants.
  2. Add a generic extension channel — the one thing an out-of-core transport genuinely needs from the daemon.
  3. Reinstate the transport at libraries/extensions/tensor-pool, on that channel, behind --features tensor-pool.

Why

#1872 — the design issue for this feature — explicitly declined to commit to an architecture:

This issue documents the gap and invites proposals. It does NOT commit to an architecture. dora has historically moved AWAY from custom shared-memory infrastructure (see #1745…). Any new transport has to clear a high bar.

Priority: not currently scheduled. Filed for visibility.

It was filed 2026-05-19, the same day #1623 (+2,717 lines) was closed for being "7× larger than the feature needed", and listed five questions a proposal had to answer before writing code. #2168 landed the same architecture 24 days later, merged 5 days after opening, answering none of them. Question 3 — "who owns pinned memory, and when is it freed? what if the producer crashes mid-write?" — is the direct ancestor of #2881, #2935 and #3015.

The problem was never the feature. It was that the feature had grown into dora's core: ~3,000 lines inside apis/python/node/src/lib.rs with 64 unsafe sites (~40 hand-parsing a 256-byte binary header at hardcoded offsets, out of memory another process writes), 950 lines of daemon lifecycle logic, pool-specific variants frozen into the wire protocol — and no CI able to exercise any of its GPU paths.

So: keep the feature, remove it from the core, and make the boundary explicit.

1. The core is clean again

apis/python/node/src/lib.rs −3,048 lines (4,064 → 1,016); unsafe count 64 → 0
binaries/daemon/src/lib.rs −973 lines: handlers, reclamation, orphan sweep
Wire protocol 3 DaemonRequest + 2 DaemonReply pool variants removed
Release workflows dora-memory-pool dropped from cargo-release.yml / release.yml, which would otherwise publish a crate that no longer exists

grep -ri 'cuda\|doradma\|pinned\|seqlock' over dora's core now returns nothing.

2. The seam

A transport outside the core still needs one thing only the daemon can provide: reclamation after a crash. A node that dies cannot withdraw the descriptor it published, so its readers keep mappings to memory nobody owns. That is #2881, and send_output cannot fix it.

So dora gained a dataflow-scoped table of opaque byte values whose lifetime the daemon brokers — and nothing else:

node.extension_store(namespace, key, value)          # bytes in
node.extension_load(namespace, key, remove=False)    # bytes out, or None
node.extension_drop(namespace, key)                  # withdraw + notify
node.drain_dropped_extension_keys(namespace)         # what went away

Same four on DoraNode for Rust nodes. Guarantees and limits: docs/extensions.md.

Deliberately generic — the variants are ExtensionStore / ExtensionLoad / ExtensionDrop, and adding a second extension needs no change to dora at all. A pool-shaped accessor would have frozen this transport's architecture into the framework, which is exactly what #1872 declined to do.

17 tests: ownership (a foreign node cannot hijack a key), namespace and dataflow scoping, the reclamation asymmetry (an owner's exit reclaims, a reader's must not), idempotent drop, per-dataflow cap.

3. The transport, opt-in

Lives at libraries/extensions/tensor-pool — beside ros2-bridge, which has the same python/ subcrate shape.

maturin develop -m apis/python/node/Cargo.toml --features tensor-pool
cargo build -p dora-daemon --features tensor-pool   # orphaned-segment reclamation

Both flags are off by default and independent. The default build neither compiles nor exposes any of it: process_pending_tensor_pool_frees becomes an empty no-op, and default cargo check on the daemon and Python crate is warning-clean.

register_tensor_pool and friends are still methods on Node, #[cfg(feature)]-gated, delegating to a Pool context struct (PyO3 cannot let another crate add #[pymethods]).

Renamed from memory-pool. dora already has an unrelated shared_memory_pool_size descriptor key (and DORA_NODE_SHM_POOL_SIZE) for the Zenoh SHM buffer pool, and "memory pool" for both was a persistent source of confusion. tensor-pool says what this one actually pools and stays accurate on the CPU path, which gpu- would not — the CPU path works without CUDA and is the only one with CI coverage.

This renames the four Python methods (register_memory_poolregister_tensor_pool, and likewise for write/read/free), the crates, the feature flags and the example env keys. That is a user-visible break for anyone on the current methods — permissible precisely because this feature sits outside the 1.0 guarantees, and better done now than after more users arrive.

Its six former daemon calls now go through the seam:

Before Now
register_pinned_memory(id, meta) extension_store("dora-tensor-pool", id, bytes)
read_pinned_memory(id, free) ×3 extension_load("dora-tensor-pool", id, remove=free)
free_pinned_memory(id) extension_drop("dora-tensor-pool", id)
drain_freed_pools() drain_dropped_extension_keys("dora-tensor-pool")

The descriptor is JSON that dora never parses. That boundary is what keeps the unsafe pointer arithmetic, the seqlock and the embedded libcudart bindings on the extension's side.

Guarantee status is stated where it gets read: the README leads with it, and it is repeated in the crate description, the module header, every #[pymethods] docstring and the .pyi stubs. Named open defects: #3015, #2935, #2890. (#2881 is fixed.)

Three latent bugs surfaced

Returning the crate to the workspace put it under -D warnings for the first time, which caught: a deprecated downcast_into, a dead initializer in the device-to-host copy path, and seqlock_begin_write with no callers at all — every write path uses begin_if_even. All fixed; the dead function is deleted.

Known gap

The smoke tests are not wired into any automated suite. They need a feature-built wheel plus torch, so as a cargo test target they broke the build. They ship as smoke-tests.rs.example with instructions instead. The extension therefore has 53 unit tests but no in-tree end-to-end coverage — a real regression against the old nightly memory-pool-smoke job, and worth tracking separately.

Validation

cargo fmt --all -- --check, cargo clippy --all -- -D warnings, cargo check --examples — clean, on both the default and --features tensor-pool builds.

cargo test --all: two failures, dora-ros2-bridge --test rmw_zenoh_pubsub::{transient_local_delivers_history_to_late_joiner, two_sessions_preserve_order_payload_and_metadata}, both Elapsed(()) on Zenoh peer discovery. They reproduce identically on an unmodified checkout — the dev container has no multicast — and this branch touches zero lines in that crate. The apis/c++ build-script failure in that environment is likewise pre-existing on main.

@trunk-io

trunk-io Bot commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

😎 Merged successfully - details.

phil-opp and others added 2 commits August 13, 2026 14:26
Extracts the pinned/CUDA memory-pool transport (#2168, #2386, #2619) into
external/dora-pool, staged for lifting into its own repository.

dora 1.0 ships no pool API: the four Python methods, the three
DaemonRequest/two DaemonReply variants, the daemon-side registry and the
node-api plumbing are all removed. No seam is left behind — #1872 declined
to commit to this architecture, so dora should not ship a socket moulded to
its shape.

The parked copy is the post-#3014 code (pool reclamation, #2881), so it
carries that fix rather than the pre-fix state.

external/dora-pool/README.md leads with a seam contract: a budget for any
future reinstatement (<200 lines of dora, no new unsafe), a table of what
may never return in-tree, and re-entry criteria — #1872's five unanswered
design questions, the open correctness bugs, and a GPU CI story.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…nsports

A transport that lives outside dora still needs one thing only the daemon can
provide: reclamation after a crash. A node that dies cannot withdraw the
descriptor it published, so its readers keep mappings to memory nobody owns
(#2881 is that failure mode with a real transport attached).

Adds a dataflow-scoped table of opaque byte values whose lifetime the daemon
brokers — store / load / drop, plus a drained notification when a key goes
away. dora never interprets the namespace, key or value.

Deliberately generic rather than shaped around any one transport: naming the
protocol variants after the memory pool would freeze that architecture into
dora, which #1872 explicitly declined to do. A second extension needs no
change here at all.

Guarantees: only the storing node may overwrite a key; entries are scoped per
dataflow and per namespace; every node that stored or read a key is notified
when it is dropped; a dropped key is reclaimed on owner exit and on dataflow
finish; dropping an absent key is a no-op so retries are safe. Bounded at 8192
entries per dataflow and 4096 pending notifications per process.

This is a control plane for descriptors, not a data plane — values are copied
through the daemon. See docs/extensions.md.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@phil-opp phil-opp changed the title refactor: park the memory-pool transport out of the tree refactor(memory-pool): move behind a generic extension seam, opt-in and outside the 1.0 guarantees Aug 13, 2026

Copy link
Copy Markdown
Collaborator Author

🤖 Automated review by Claude — fully automated review; no human has verified these findings.

No issues found.

I checked the new generic extension channel and its plumbing (the rest of the diff is code movement into libraries/extensions/memory-pool and the removal of the in-tree pool API):

  • binaries/daemon/src/extension_table.rs — store/load/drop/reclaim_owner/reclaim_dataflow, owner-only overwrite, and the per-dataflow cap all look correct and are well covered by unit tests (including the overwrite-at-cap edge case).
  • apis/rust/node/src/event_stream/extensions.rs — the bounded, namespace-scoped drop-notification queue is sound.
  • libraries/message/src/{node_to_daemon,daemon_to_node}.rs — every match arm (payload_len, wants_reply, expects_tcp_reply, encode_size_hint) is updated consistently for the three new variants.
  • Daemon handlers in lib.rsExtensionLoad { remove: true } only broadcasts a drop on a hit, ExtensionDrop is idempotent, and inc_pending accounting matches the prior pool path.
  • cuda.py — the remaining functions don't reference the removed _DTYPE_MAP/_TORCH_TO_NUMPY_DTYPE_MAP, so no dangling references.

One observation, not a bug: unlike the removed notify_memory_pool_freed (which excluded the initiator), notify_extension_dropped now also delivers ExtensionDropped to the node that triggered the drop / load-with-remove, since it's in touched_by. That's consumed out-of-band by the event-stream thread and left for the extension to drain, and the eager owner-exit reclamation deliberately replaces the old reachability logic (downstream_closure/remote_edges) — both are documented as intentional, with the pool extension marked opt-in and outside the 1.0 guarantees.


Generated by Claude Code

Keeps the feature and its Python API surface, but out of dora's core: it
lives at libraries/extensions/tensor-pool (beside ros2-bridge, same python/
subcrate shape) and reaches dora only through the generic extension channel.

Opt-in and off by default, two independent flags:

  maturin develop -m apis/python/node/Cargo.toml --features tensor-pool
  cargo build -p dora-daemon --features tensor-pool

A default build neither compiles nor exposes it — the receive-path drain
becomes an empty no-op and dora's core keeps zero pool vocabulary.

NOT covered by the 1.0 compatibility guarantees, stated in the README, the
crate description, the module header, every pymethod docstring and the .pyi
stubs, along with the open defects (#3015, #2935, #2890).

Named tensor-pool, not memory-pool: dora already has an unrelated
shared_memory_pool_size descriptor key (and DORA_NODE_SHM_POOL_SIZE) for the
Zenoh SHM buffer pool, and one name for both was a persistent source of
confusion. 'gpu-' would have been the other obvious fix but is inaccurate —
the CPU path works without CUDA and is the only one with CI coverage. This
renames the four Python methods (register_memory_pool -> register_tensor_pool
and likewise for write/read/free), the crates, the feature flags and the
example env keys: a user-visible break, permissible because the feature sits
outside the 1.0 guarantees, and cheaper now than later.

The six former daemon calls now go through extension_store / extension_load /
extension_drop / drain_dropped_extension_keys, with the descriptor encoded as
JSON dora never parses (python/src/seam.rs). The unsafe pointer arithmetic,
the seqlock and the embedded libcudart bindings stay on the extension's side.

Returning the crate to the workspace put it under -D warnings for the first
time, which surfaced three latent bugs: a deprecated downcast_into, a dead
initializer in the device-to-host copy path, and seqlock_begin_write with no
callers at all (every write path uses begin_if_even). All fixed.

Known gap: the smoke tests need a feature-built wheel plus torch, so they
ship as smoke-tests.rs.example rather than a cargo target that cannot
compile. The extension has 53 unit tests but no in-tree end-to-end coverage.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@phil-opp phil-opp changed the title refactor(memory-pool): move behind a generic extension seam, opt-in and outside the 1.0 guarantees refactor(tensor-pool): move behind a generic extension seam, opt-in and outside the 1.0 guarantees Aug 13, 2026
@trunk-io
trunk-io Bot merged commit 0442564 into main Aug 13, 2026
16 checks passed
@trunk-io
trunk-io Bot deleted the pool-extract branch August 13, 2026 16:14
tang-canran added a commit to tang-canran/dora that referenced this pull request Aug 14, 2026
…sor-pool

Port the pre-tensor-pool cross-machine data plane to the upstream
tensor-pool extension seam (dora-rs#3152 rework):

- Daemon holds a TensorPoolManager for the cross_pools table (mirror
  tracking); main pool table stays node-side per the extension model.
- Full data plane restored: zenoh relay (MemoryPoolWrite/RegisterPool/
  RegisterPoolAck/FreePool/MemoryPoolWriteAck), mirror create/write under
  the DORADMA seqlock, direct-TCP listener + persistent connections,
  commit-ack reply withholding with seq matching, per-(dataflow,pool)
  dual-path write locks, orphan mirror sweep, per-dataflow state drain.
- Cross-machine methods moved out of the tensor-pool test module into
  the production impl (they were unreachable behind cfg(test)).
- Auth: DORA_MEMORY_POOL_AUTH_TOKEN enables a token handshake on every
  direct-TCP connection (send + verify sides); rejected peers are
  dropped before any frame touches a mirror, the origin degrades to
  the zenoh relay. Token injected into spawned nodes.
- Opt-in: the whole data plane is gated behind
  DORA_MEMORY_POOL_CROSS_MACHINE=1 (listener, mirroring, register);
  disabled daemons reject cross-machine registration with a clear error.
- Bind config: DORA_MEMORY_POOL_DATA_BIND for the data listener
  (default 0.0.0.0); existing PORT/ADDR envs unchanged.
- FreePinnedMemory daemon request restored as the cross-machine release
  entry (unregister + targeted FreePool + mirror unlink).
- postcard serialization (upstream codec) replaces bincode.

Tests: 13 cross-pool tests (incl. new auth handshake accept/reject) +
full daemon suite (228) green; clippy clean on affected crates.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
tang-canran added a commit to tang-canran/dora that referenced this pull request Aug 14, 2026
- create_cross_pool_shmem carries 8 args (dataflow/machine/pool/size/
  dtype/shape/device/sender_shmem); allow the lint like the sibling
  run_inner_with_builds helpers.
- .unwrap-budget: merge db1788f pulled in upstream dora-rs#3152 tensor-pool,
  whose python/src/transport.rs:978 doc comment contains the literal
  `.unwrap()` string, which the budget script counts (no comment
  stripping). No production unwrap was added; upstream main's own budget
  is 163, so 161 keeps us stricter while staying green.

Verified: qa-unwrap 161/161, clippy -D warnings clean on daemon.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant